#!/bin/bash

# 检查显卡型号
if nvidia-smi > /dev/null 2>&1; then
    gpu_models=$(nvidia-smi --query-gpu=name --format=csv,noheader)
    gpu_count=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
elif rocm-smi > /dev/null 2>&1; then
    rocminfo_output=$(rocminfo)
    gpu_models=$(echo "$rocminfo_output" | grep "Marketing Name" | grep "Radeon" | awk -F': ' '{print $2}' | sed 's/^[ \t]*//;s/[ \t]*$//')
    gpu_count=$(lspci | grep -i vga | wc -l)
else
    echo "未检测到NVIDIA或AMD显卡，脚本退出。"
    exit 1
fi

# 输出结果
echo "显卡型号: $gpu_models"
echo "显卡数量: $gpu_count"
